import numpy as np
import matplotlib.pyplot as plt
from brian2 import *
import networkx as nx

# 设置Brian2的单位
defaultclock.dt = 0.1*ms

# 创建神经元群
N = 100  # 神经元数量
tau = 10*ms  # 时间常数
eqs = '''
dv/dt = (v0 - v)/tau : volt
v0 : volt
'''

# 创建神经元群组 - 随编号递增的电压
neurons1 = NeuronGroup(N, eqs,
                     threshold='v>15*mV',
                     reset='v=0*mV',
                     method='exact')
neurons1.v0 = '20*mV * i / (N-1)'

# 创建神经元群组 - 随机电压
neurons2 = NeuronGroup(N, eqs,
                     threshold='v>15*mV',
                     reset='v=0*mV',
                     method='exact')
neurons2.v0 = '20*mV * rand()'

# 创建突触连接 - 群组1
S1 = Synapses(neurons1, neurons1, 'w : volt', on_pre='v_post += w')
S1.connect(condition='i!=j', p=0.1)  # 10%的连接概率
S1.w = '0.5*mV'

# 创建突触连接 - 群组2
S2 = Synapses(neurons2, neurons2, 'w : volt', on_pre='v_post += w')
S2.connect(condition='i!=j', p=0.1)  # 10%的连接概率
S2.w = '0.5*mV'

# 记录数据 - 群组1
M1 = StateMonitor(neurons1, 'v', record=True)
spikemon1 = SpikeMonitor(neurons1)

# 记录数据 - 群组2
M2 = StateMonitor(neurons2, 'v', record=True)
spikemon2 = SpikeMonitor(neurons2)

# 运行模拟
duration = 500*ms
run(duration)

# 创建图形
plt.figure(figsize=(20, 15))

# 1. 膜电位随时间变化 - 群组1
plt.subplot(421)
for i in range(0, N, 10):
    plt.plot(M1.t/ms, M1.v[i]/mV)
plt.xlabel('Time (ms)')
plt.ylabel('v (mV)')
plt.title('Membrane Potential over Time (Sequential)')

# 1. 膜电位随时间变化 - 群组2
plt.subplot(422)
for i in range(0, N, 10):
    plt.plot(M2.t/ms, M2.v[i]/mV)
plt.xlabel('Time (ms)')
plt.ylabel('v (mV)')
plt.title('Membrane Potential over Time (Random)')

# 2. 脉冲发放图 - 群组1
plt.subplot(423)
plt.plot(spikemon1.t/ms, spikemon1.i, '.k')
plt.xlabel('Time (ms)')
plt.ylabel('Neuron index')
plt.title('Spike Raster Plot (Sequential)')

# 2. 脉冲发放图 - 群组2
plt.subplot(424)
plt.plot(spikemon2.t/ms, spikemon2.i, '.k')
plt.xlabel('Time (ms)')
plt.ylabel('Neuron index')
plt.title('Spike Raster Plot (Random)')

# 3. 网络连接可视化 - 群组1
plt.subplot(425)
G1 = nx.DiGraph()
for i in range(N):
    G1.add_node(i)
for i, j in zip(S1.i, S1.j):
    G1.add_edge(i, j)

pos1 = nx.spring_layout(G1)
nx.draw(G1, pos1, node_size=20, node_color='b', 
        with_labels=False, arrows=False)
plt.title('Network Connectivity (Sequential)')

# 3. 网络连接可视化 - 群组2
plt.subplot(426)
G2 = nx.DiGraph()
for i in range(N):
    G2.add_node(i)
for i, j in zip(S2.i, S2.j):
    G2.add_edge(i, j)

pos2 = nx.spring_layout(G2)
nx.draw(G2, pos2, node_size=20, node_color='r', 
        with_labels=False, arrows=False)
plt.title('Network Connectivity (Random)')

# 4. 平均发放率 - 群组1
plt.subplot(427)
spike_counts1 = np.histogram(spikemon1.i, bins=N)[0]
plt.bar(range(N), spike_counts1)
plt.xlabel('Neuron index')
plt.ylabel('Spike count')
plt.title('Firing Rate Distribution (Sequential)')

# 4. 平均发放率 - 群组2
plt.subplot(428)
spike_counts2 = np.histogram(spikemon2.i, bins=N)[0]
plt.bar(range(N), spike_counts2)
plt.xlabel('Neuron index')
plt.ylabel('Spike count')
plt.title('Firing Rate Distribution (Random)')

plt.tight_layout()
plt.savefig('snn_visualization.png')
plt.close()

print("模拟完成！可视化结果已保存为 'snn_visualization.png'")